import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 数据集路径
SIZE=256
BATCH_SIZE=16
train_dir = r"animals10\train"
test_dir = r"animals10\test"

def load():
    # 数据预处理
    # 定义训练集和测试集的转换操作
    train_transforms = transforms.Compose([
        transforms.Resize((SIZE, SIZE)),  # 调整图片大小
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.ToTensor(),  # 转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
    ])

    test_transforms = transforms.Compose([
        transforms.Resize((SIZE,SIZE)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 加载数据集
    train_dataset = datasets.ImageFolder(root=train_dir, transform=train_transforms)
    test_dataset = datasets.ImageFolder(root=test_dir, transform=test_transforms)

    # 创建 DataLoader
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)
    
    return train_loader, test_loader, len(train_dataset), len(test_dataset)

    # 打印数据集信息
    print(f"Train dataset size: {len(train_dataset)}")
    print(f"Test dataset size: {len(test_dataset)}")
    print(f"Classes: {train_dataset.classes}")